from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

app = FastAPI()

class Item(BaseModel):
    text: str

class Response(BaseModel):
    translated_text: str

class R(BaseModel):
    code: int
    msg: str
    data: object

    def ok():
        return R(code=200, msg='success', data=None)

    def ok(data:object):
        return R(code=200, msg='success', data=data)

    def fail(msg:str):
        return R(code=500, msg=msg, data=None)

    def fail(msg:str, data:object):
        return R(code=500, msg=msg, data=data)

# 加载模型
# 根据输入模型选择加载什么模型，使用虚拟环境为translation,token 最大长度是512
en2zhModel = None
en2zhTokenizer = None
zh2enModel = None
zh2enTokenizer = None
device = 'cpu'

input_model = 'opus' # 可选 t5,opus

if input_model == 't5':
    from t5_run import getModelAndTokenizer
    en2zhModel, en2zhTokenizer = getModelAndTokenizer()
    zh2enModel, zh2enTokenizer = getModelAndTokenizer()
    print('load t5 model')
elif input_model == 'opus':
    from opus_en2zh import getModelAndTokenizer as en2zhGetModelAndTokenizer
    from opus_zh2en import getModelAndTokenizer as zh2enGetModelAndTokenizer
    en2zhModel,en2zhTokenizer = en2zhGetModelAndTokenizer()
    zh2enModel,zh2enTokenizer = zh2enGetModelAndTokenizer()
    print('load opus model')

port = 17737


@app.post("/translate/zh2en")
async def translate_zh2en(item: Item):
    try:
        # print("翻译前的文本长度: ",len(item.text))
        input_ids = zh2enTokenizer(item.text, return_tensors="pt")

        generated_tokens = zh2enModel.generate(**input_ids.to(device))

        result = zh2enTokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        res = Response(translated_text=result[0])
        return R.ok(res)
    except Exception as e:
        return R.fail(str(e))

@app.post("/translate/en2zh")
async def translate_en2zh(item: Item):
    try:
        # print("翻译前的文本长度: ",len(item.text))
        input_ids = en2zhTokenizer(item.text, return_tensors="pt")

        generated_tokens = en2zhModel.generate(**input_ids.to(device))

        result = en2zhTokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        res = Response(translated_text=result[0])
        return R.ok(res)
    except Exception as e:
        return R.fail(str(e))

# 后台运行 nohup python translation_main.py 2>&1 &
if __name__ == '__main__':
    # 设置运行端口
    uvicorn.run(app, host="0.0.0.0", port=port)